from hep.aida import *

from java.awt import Color
from java.awt import Font
from java.lang import *
from java.util import *
from java.io import *

from jhplot.jadraw import *
from jhplot  import *
from jplot  import *

from fr.in2p3.ipnl.aidautilities import *
from fr.in2p3.ipnl.oflap.scint import *
from jscinttools.utilities.jhplotutilities import *

Locale.setDefault(Locale.ENGLISH)
System.getProperties().setProperty("hep.aida.BatchAnalysisFactory","hep.aida.ref.BatchAnalysisFactory")

showFigures = False
writeFigures = not showFigures

useDiameterRatherThanArea = False

""" Define colors and fonts"""
figType = [".png", ".eps"]
colors = [Color.blue, Color.red, Color(0, 170, 0), Color.black, Color.yellow]
axisLabelFont = Font(Font.SERIF, Font.PLAIN, 25)
titleFont = Font(Font.SERIF, Font.PLAIN, 25)
legendFont = Font(Font.SERIF, Font.PLAIN, 20)
ticFont = Font(Font.SERIF, Font.PLAIN, 20)
statBoxFont = Font(Font.SERIF, Font.PLAIN, 20)


""" Define the Path and files names"""
#aidaPathL = "/home/marc-antoine/Documents/BGO/QFin"
#aidaFileL = "ScintillationAnalysis_BGO"
aidaPathQ = "/home/marc-antoine/Documents/BGO/QFin/QueensData/queenstools"
aidaFileQ = "ScintillationAnalysis_BGO_QFile"
outPath = "/home/marc-antoine/Documents/Dropbox/BGO"
#afrL = AidaFileReader(aidaPathL, aidaFileL)
#afrQ = AidaFileReader(aidaPathQ, aidaFileQ)


""" Analysis Factories """
af = IAnalysisFactory.create()
treeName = "KinComp"
mainTree = af.createTreeFactory().createTree(outPath + File.separator + treeName + ".aida", "xml", ITreeFactory.CREATE)
dpsf = af.createDataPointSetFactory(mainTree)
aUtil = AidaUtil(mainTree)

"""
###############################################################
# Plotting Part
###############################################################
"""
""" Create the plottable JHep objects from AIDA IObjects"""
shortDPSQ = []
mainDPSQ = []
longDPSQ = []
shortQ = []
mainQ = []
longQ = []
shortContQ = []
mainContQ = []
longContQ = []

treeQ = AidaFileReader(aidaPathQ, aidaFileQ).tree
shQ = GrabIObject(treeQ,"/Kinetic/Decay_Q_Comp#1").dps
mnQ = GrabIObject(treeQ,"/Kinetic/Decay_Q_Comp#2").dps
lgQ = GrabIObject(treeQ,"/Kinetic/Decay_Q_Comp#3").dps
aUtil.sortUp(shQ)
aUtil.sortUp(mnQ)
aUtil.sortUp(lgQ)


shNumQ = GrabIObject(treeQ,"/Light_Yield/LY_Q_Comp#1").dps
mnNumQ = GrabIObject(treeQ,"/Light_Yield/LY_Q_Comp#2").dps
lgNumQ = GrabIObject(treeQ,"/Light_Yield/LY_Q_Comp#3").dps

aUtil.sortUp(shNumQ)
aUtil.sortUp(mnNumQ)
aUtil.sortUp(lgNumQ)

for i in range(shQ.size()):
	if (shQ.point(i).coordinate(1).value() != Double.NaN) and (shNumQ.point(i).coordinate(1).value() != Double.NaN):
		shortDPSQ.append(dpsf.create("shQPt"+String.valueOf(i), 2))
		shortDPSQ[i].addPoint();
		shortDPSQ[i].point(0).coordinate(0).setValue(shQ.point(i).coordinate(0).value())
		shortDPSQ[i].point(0).coordinate(1).setValue(shQ.point(i).coordinate(1).value())
		shortDPSQ[i].point(0).coordinate(1).setErrorMinus(shQ.point(i).coordinate(1).errorMinus())
		shortDPSQ[i].point(0).coordinate(1).setErrorPlus(shQ.point(i).coordinate(1).errorPlus())
		
		shortQ.append(P1D(shortDPSQ[i]))
		shortQ[i].setSymbol(4)
		shortQ[i].color = colors[2]
		shortQ[i].errY = True
		shortQ[i].errX = False
		shortQ[i].errColor = shortQ[i].getColor()
		jpl = shortQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl.setSymbolSize(shNumQ.point(i).coordinate(1).value() * 10)
		else:
			jpl.setSymbolSize(Math.sqrt(shNumQ.point(i).coordinate(1).value()) * 18)
		shortQ[i].setDrawOption(jpl)
		
		shortContQ.append(P1D(shortDPSQ[i]))
		shortContQ[i].setSymbol(4)
		shortContQ[i].color = colors[3]
		shortContQ[i].errY = False
		shortContQ[i].errX = False
		jpl2 = shortContQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl2.setSymbolSize((shNumQ.point(i).coordinate(1).value() * 10) + 2)
		else:
			jpl2.setSymbolSize((Math.sqrt(shNumQ.point(i).coordinate(1).value()) * 18) + 2)
		#jpl2.setSymbolSize((shNumQ.point(i).coordinate(1).value() * 10) + (shNumQ.point(i).coordinate(1).errorMinus() * 20) + 2)
		shortContQ[i].setDrawOption(jpl2)
	
for i in range(mnQ.size()):
	if (mnQ.point(i).coordinate(1).value() != Double.NaN) and (mnNumQ.point(i).coordinate(1).value() != Double.NaN):
		mainDPSQ.append(dpsf.create("mnQPt"+String.valueOf(i), 2))
		mainDPSQ[i].addPoint();
		mainDPSQ[i].point(0).coordinate(0).setValue(mnQ.point(i).coordinate(0).value())
		mainDPSQ[i].point(0).coordinate(1).setValue(mnQ.point(i).coordinate(1).value())
		mainDPSQ[i].point(0).coordinate(1).setErrorMinus(mnQ.point(i).coordinate(1).errorMinus())
		mainDPSQ[i].point(0).coordinate(1).setErrorPlus(mnQ.point(i).coordinate(1).errorPlus())
		
		mainQ.append(P1D(mainDPSQ[i]))
		mainQ[i].setSymbol(4)
		mainQ[i].color = colors[0]
		mainQ[i].errY = True
		mainQ[i].errX = False
		mainQ[i].errColor = mainQ[i].getColor()
		jpl = mainQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl.setSymbolSize(mnNumQ.point(i).coordinate(1).value() * 10)
		else:
			jpl.setSymbolSize(Math.sqrt(mnNumQ.point(i).coordinate(1).value()) * 18)
		mainQ[i].setDrawOption(jpl)
		
		mainContQ.append(P1D(mainDPSQ[i]))
		mainContQ[i].setSymbol(4)
		mainContQ[i].color = colors[3]
		mainContQ[i].errY = False
		mainContQ[i].errX = False
		jpl2 = mainContQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl2.setSymbolSize((mnNumQ.point(i).coordinate(1).value() * 10) + 4)
		else:
			jpl2.setSymbolSize((Math.sqrt(mnNumQ.point(i).coordinate(1).value()) * 18) + 4)
		#jpl2.setSymbolSize((mnNumQ.point(i).coordinate(1).value() * 10) + (mnNumQ.point(i).coordinate(1).errorMinus() * 20) + 2)
		mainContQ[i].setDrawOption(jpl2)
	
for i in range(lgQ.size()):
	if (lgQ.point(i).coordinate(1).value() != Double.NaN) and (lgNumQ.point(i).coordinate(1).value() != Double.NaN):
		longDPSQ.append(dpsf.create("lgQPt"+String.valueOf(i), 2))
		longDPSQ[i].addPoint();
		longDPSQ[i].point(0).coordinate(0).setValue(lgQ.point(i).coordinate(0).value())
		longDPSQ[i].point(0).coordinate(1).setValue(lgQ.point(i).coordinate(1).value())
		longDPSQ[i].point(0).coordinate(1).setErrorMinus(lgQ.point(i).coordinate(1).errorMinus())
		longDPSQ[i].point(0).coordinate(1).setErrorPlus(lgQ.point(i).coordinate(1).errorPlus())
		
		longQ.append(P1D(longDPSQ[i]))
		longQ[i].setSymbol(4)
		longQ[i].color = colors[1]
		longQ[i].errY = True
		longQ[i].errX = False
		longQ[i].errColor = longQ[i].getColor()
		jpl = longQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl.setSymbolSize(lgNumQ.point(i).coordinate(1).value() * 10)
		else:
			jpl.setSymbolSize(Math.sqrt(lgNumQ.point(i).coordinate(1).value()) * 18)
		longQ[i].setDrawOption(jpl)
		
		longContQ.append(P1D(longDPSQ[i]))
		longContQ[i].setSymbol(4)
		longContQ[i].color = colors[3]
		longContQ[i].errY = False
		longContQ[i].errX = False
		jpl2 = longContQ[i].getDrawOption()
		if useDiameterRatherThanArea:
			jpl2.setSymbolSize((lgNumQ.point(i).coordinate(1).value() * 10) + 4)
		else:
			jpl2.setSymbolSize((Math.sqrt(lgNumQ.point(i).coordinate(1).value()) * 18) + 4)
		#jpl2.setSymbolSize((lgNumQ.point(i).coordinate(1).value() * 10) + (lgNumQ.point(i).coordinate(1).errorMinus() * 20) + 2)
		longContQ[i].setDrawOption(jpl2)

""" Create the Keys for the Legend """
shKeyQ = JaKey(shortQ[0].getSymbol(), "Fast Comp", 0.20, 0.80, "NDC")
shKeyQ.color = shortQ[0].getColor()
shKeyQ.keyColor = shortQ[0].getColor()
shKeyQ.penWidth = 2
shKeyQ.font = legendFont

mnKeyQ = JaKey(mainQ[0].getSymbol(), "Main Comp", 0.20, 0.75, "NDC")
mnKeyQ.color = mainQ[0].getColor()
mnKeyQ.keyColor = mainQ[0].getColor()
mnKeyQ.penWidth = 2
mnKeyQ.font = legendFont

lgKeyQ = JaKey(longQ[0].getSymbol(), "Afterglow", 0.20, 0.70, "NDC")
lgKeyQ.color = longQ[0].getColor()
lgKeyQ.keyColor = longQ[0].getColor()
lgKeyQ.penWidth = 2
lgKeyQ.font = legendFont

#txt = JaText("Diameter &prop; LY", 0.60, 0.2, "NDC")
lab = ""
if useDiameterRatherThanArea:
	lab = "Diameter &prop; LY"
else:
	lab = "Area &prop; LY"
txt = JaText(lab, 0.55, 0.2, "NDC")
txt.font = statBoxFont
txt.color = colors[0]

txt2 = JaText("Border &prop; LY error", 0.55, 0.25, "NDC")
txt2.font = statBoxFont
txt2.color = colors[0]

""" Create the Plotter"""
plotter = HPlotJa("KinComp", 600, 600, 1, 1, True)
plotter.setGTitle("Scintillation Structure of BGO", titleFont)
plotter.antiAlias = False
if showFigures:
	plotter.visible(True)
else:
	plotter.visible(False)
	
plotter.setPad(1, 1, 0.125, 0.125, 0.75, 0.75)
plotter.setRange(1, 1000, 0.01, 1000)
plotter.setLogScale(0, True)
plotter.setLogScale(1, True)
plotter.ticFont = ticFont
plotter.showStatBox = False
plotter.showKey(False)
plotter.nameX = "Temperature (K)"
plotter.nameY = "Decay Time (&micro;s)"

for i in range(len(mainQ)-1,-1,-1):
	plotter.draw(mainContQ[i])
	plotter.draw(mainQ[i])

for i in range(len(shortQ)-1,-1,-1):
	plotter.draw(shortContQ[i])
	plotter.draw(shortQ[i])
	
for i in range(len(longQ)-1,-1,-1):
	plotter.draw(longContQ[i])
	plotter.draw(longQ[i])
	
plotter.draw(shKeyQ)
plotter.draw(mnKeyQ)
plotter.draw(lgKeyQ)
plotter.draw(txt)
#plotter.draw(txt2)
	
""" Modifies the JaAxes and the axes labels"""
labX = plotter.getLabelX()[0][0]
labX.font = axisLabelFont
HPlotJaUtil.setLabelLocationAtCenter(0, labX, plotter)
labY = plotter.getLabelY()[0][0]
labY.font = axisLabelFont
#HPlotJaUtil.setLabelLocationAtCenter(1, labY, plotter, -0.01)
HPlotJaUtil.setLabelLocationAtCenter(1, labY, plotter, -0.01)


theAxes = plotter.getPad()
theAxes.setLabelSpace(0, 0.0)
#theAxes.setLabelSpace(1, -0.025)
theAxes.setLabelSpace(1, -0.005)
HPlotJaUtil.setTickLabelOneToPower(0, theAxes)
HPlotJaUtil.setTickLabelOneToPower(1, theAxes)

plotter.update()

""" Saves or not the plot"""
if writeFigures:
	for i in range(len(figType)):
		outFileName = "KinCompV5NerArea" + figType[i]
		plotter.export(outPath + File.separator + outFileName)
		print("Writting File : " + outPath + File.separator + outFileName)
		
plotter.update()
print("That's all folks !")

